import tensorflow as tf
from tensorflow.keras import metrics


@tf.function
def ks(y_true, y_pred):
    """
    ks：金融风控指标（二分类问题）

    ks = max(TPR - FPR)
    TPR = TP / (TP + FN)：正样本的累积分布曲线(CDF)
    FPR = FP / (FP + TN)：负样本的累积分布曲线(CDF)

    函数式评估指标

    :param y_true:
    :param y_pred:
    :return:
    """
    # 转为一维数组
    y_true = tf.reshape(y_true, (-1, ))
    y_pred = tf.reshape(y_pred, (-1, ))

    # 打平后一维数组的长度
    length = tf.shape(y_true)[0]
    # 最大的前k个
    max_idx = tf.math.top_k(input=y_pred, k=length, sorted=False)

    y_true_sorted = tf.gather(y_true, max_idx.indices)
    # y_pred_sorted = tf.gather(y_pred, max_idx.indices)

    """
    tf.truediv: 除法， x / y
    tf.reduce_sum: 指定维度求和，若维度为None，则对所有元素求和
    """
    tpr = tf.truediv(x=tf.cumsum(y_true_sorted), y=tf.reduce_sum(y_true_sorted))
    fpr = tf.truediv(x=tf.cumsum(1 - y_true_sorted), y=tf.reduce_sum(1 - y_true_sorted))

    return tf.reduce_max(tf.abs(tpr - fpr))


class KS(metrics.Metric):
    """
    ks：金融风控指标（二分类问题）

    ks = max(TPR - FPR)
    TPR = TP / (TP + FN)：正样本的累积分布曲线(CDF)
    FPR = FP / (FP + TN)：负样本的累积分布曲线(CDF)

    类式评估指标
    """
    def __init__(self, name='ks', **kwargs):
        super(KS, self).__init__(name=name, **kwargs)
        self.true_positives = self.add_weight(name='tp', shape=(101, ), initializer='zeros')
        self.false_positives = self.add_weight(name='fp', shape=(101, ), initializer='zeros')

    @tf.function
    def update_state(self, y_true, y_pred):
        """
        更新相关中间变量的状态
        :param y_true:
        :param y_pred:
        :return:
        """
        y_true = tf.cast(tf.reshape(y_true, (-1, )), tf.bool)
        # 百分数转换
        y_pred = tf.cast(100 * tf.reshape(y_pred, (-1, )), tf.int32)
        # 长度
        length = tf.shape(y_true)[0]
        for i in tf.range(length):
            # 权重更新
            if y_true[i]:
                self.true_positives[y_pred[i]].assign(self.true_positives[y_pred[i]] + 1.0)
            else:
                self.false_positives[y_pred[i]].assign(self.false_positives[y_pred[i]] + 1.0)
        return (self.true_positives, self.false_positives)

    @tf.function
    def result(self):
        """
        输出最终指标结果
        :return:
        """
        tpr = tf.truediv(tf.cumsum(self.true_positives), tf.reduce_sum(self.true_positives))
        fpr = tf.truediv(tf.cumsum(self.false_positives), tf.reduce_sum(self.false_positives))
        return tf.reduce_max(tf.abs(tpr - fpr))


def run():
    y_true = tf.constant([[1], [1], [1], [0], [1], [1], [1], [0], [0], [0], [1], [0], [1], [0]])
    y_pred = tf.constant([[0.6], [0.1], [0.4], [0.5], [0.7], [0.7], [0.7], [0.4], [0.4], [0.5],
                          [0.8], [0.3], [0.5], [0.3]])
    # 函数式
    # ks_value = ks(y_true=y_true, y_pred=y_pred)

    # 类式
    ks = KS()
    ks.update_state(y_true=y_true, y_pred=y_pred)
    ks_value = ks.result()

    tf.print(ks_value)


if __name__ == '__main__':
    run()
